#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys
import json
import time
import threading
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
                            QLabel, QLineEdit, QPushButton, QTextEdit, 
                            QGroupBox, QGridLayout, QCheckBox, 
                            QMessageBox)
from PyQt5.QtCore import Qt, QTimer, pyqtSignal, QThread
import websocket
import logging

class WebSocketThread(QThread):
    """WebSocket客户端线程类"""
    connectionChanged = pyqtSignal(bool, str)  # 连接状态变化信号
    
    def __init__(self):
        super().__init__()
        self.ws = None
        self.is_running = False
        self.is_connecting = False
        self.host = "localhost"
        self.port = 8765
        self.reconnect_attempts = 0
        self.max_reconnect_attempts = 10
        self.reconnect_timer = None
        self.message_queue = []
    
    def run(self):
        self.is_running = True
        self.connect_websocket()
        
        while self.is_running:
            time.sleep(0.1)
    
    def connect_websocket(self):
        if self.ws and self.ws.sock and self.ws.sock.connected:
            logging.info("WebSocket已连接，跳过连接")
            return
        
        if self.is_connecting:
            logging.info("WebSocket正在连接中，跳过连接")
            return
        
        try:
            self.is_connecting = True
            self.connectionChanged.emit(False, "正在连接...")
            
            if self.ws:
                self.ws.close()
                self.ws = None
            
            ws_url = f"ws://{self.host}:{self.port}"
            
            self.ws = websocket.WebSocketApp(
                ws_url,
                on_open=self.on_open,
                on_message=self.on_message,
                on_error=self.on_error,
                on_close=self.on_close
            )
            
            threading.Thread(target=self.ws.run_forever).start()
            
        except Exception as e:
            logging.error(f"连接WebSocket服务器失败: {str(e)}")
            self.is_connecting = False
            self.connectionChanged.emit(False, f"连接失败: {str(e)}")
            self.try_reconnect()
    
    def on_open(self, ws):
        logging.info("WebSocket连接成功")
        self.is_connecting = False
        self.reconnect_attempts = 0
        self.connectionChanged.emit(True, "已连接")
        
        if self.message_queue:
            for message in self.message_queue:
                self.send_message(message)
            self.message_queue.clear()
        
        QTimer.singleShot(0, self.start_heartbeat)
    
    def on_message(self, ws, message):
        logging.debug(f"收到消息: {message}")
    
    def on_error(self, ws, error):
        logging.error(f"WebSocket错误: {str(error)}")
        self.connectionChanged.emit(False, f"连接错误: {str(error)}")
    
    def on_close(self, ws, close_status_code, close_msg):
        logging.info("WebSocket连接已关闭")
        self.is_connecting = False
        self.connectionChanged.emit(False, "连接已关闭")
        self.try_reconnect()
    
    def try_reconnect(self):
        if not self.is_running:
            return
            
        if self.reconnect_attempts >= self.max_reconnect_attempts:
            logging.error("达到最大重连次数，停止重连")
            return
        
        self.reconnect_attempts += 1
        delay = min(1000 * (2 ** self.reconnect_attempts), 30000)
        
        logging.info(f"将在{delay/1000}秒后尝试第{self.reconnect_attempts}次重连")
        self.connectionChanged.emit(False, f"将在{delay/1000}秒后重连 ({self.reconnect_attempts}/{self.max_reconnect_attempts})")
        
        if self.reconnect_timer:
            self.reconnect_timer.stop()
            
        self.reconnect_timer = QTimer()
        self.reconnect_timer.setSingleShot(True)
        self.reconnect_timer.timeout.connect(self.connect_websocket)
        self.reconnect_timer.start(delay)
    
    def send_message(self, message):
        if not isinstance(message, (dict, str)):
            logging.error("无效的消息格式")
            return False
        
        if isinstance(message, dict):
            message = json.dumps(message)
            
        if self.ws and self.ws.sock and self.ws.sock.connected:
            try:
                self.ws.send(message)
                logging.info(f"消息已发送: {message[:100]}...")
                return True
            except Exception as e:
                logging.error(f"发送消息失败: {str(e)}")
                return False
        else:
            logging.warning("WebSocket未连接，将消息加入队列")
            if isinstance(message, str):
                self.message_queue.append(message)
            return False
    
    def start_heartbeat(self):
        heartbeat_timer = QTimer()
        heartbeat_timer.timeout.connect(self.send_heartbeat)
        heartbeat_timer.start(15000)
        
        self.heartbeat_timer = heartbeat_timer
    
    def send_heartbeat(self):
        if self.ws and self.ws.sock and self.ws.sock.connected:
            try:
                self.ws.send(json.dumps({
                    "type": "heartbeat", 
                    "timestamp": int(time.time() * 1000)
                }))
            except Exception as e:
                logging.error(f"发送心跳包失败: {str(e)}")
                if self.heartbeat_timer:
                    self.heartbeat_timer.stop()
                self.ws.close()
    
    def set_server(self, host, port):
        self.host = host
        self.port = port
    
    def stop(self):
        self.is_running = False
        
        if self.heartbeat_timer:
            self.heartbeat_timer.stop()
            
        if self.reconnect_timer:
            self.reconnect_timer.stop()
            
        if self.ws:
            self.ws.close()


class MainWindow(QMainWindow):
    """主窗口类"""
    def __init__(self):
        super().__init__()
        self.ws_thread = None
        self.data_to_send = []
        self.current_data_index = 0
        self.is_sending = False
        self.send_timer = QTimer()
        self.send_timer.timeout.connect(self.send_next_data)
        self.setup_ui()
        self.init_websocket()
        
    def setup_ui(self):
        # 窗口基本属性
        self.setWindowTitle("WebSocket发送工具")
        self.resize(600, 400)
        
        # 创建中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)
        
        # === 顶部连接设置 ===
        connection_group = QGroupBox("连接设置")
        connection_layout = QGridLayout(connection_group)
        
        # 主机设置
        host_label = QLabel("主机:")
        self.host_input = QLineEdit("localhost")
        connection_layout.addWidget(host_label, 0, 0)
        connection_layout.addWidget(self.host_input, 0, 1)
        
        # 端口设置
        port_label = QLabel("端口:")
        self.port_input = QLineEdit("8765")
        self.port_input.setMaximumWidth(100)
        connection_layout.addWidget(port_label, 0, 2)
        connection_layout.addWidget(self.port_input, 0, 3)
        
        # 连接按钮
        self.connect_btn = QPushButton("连接")
        self.connect_btn.clicked.connect(self.toggle_connection)
        connection_layout.addWidget(self.connect_btn, 0, 4)
        
        # 状态显示
        status_label = QLabel("状态:")
        self.status_text = QLabel("未连接")
        self.status_text.setStyleSheet("color: red;")
        connection_layout.addWidget(status_label, 0, 5)
        connection_layout.addWidget(self.status_text, 0, 6)
        
        # 添加到主布局
        main_layout.addWidget(connection_group)
        
        # === 发送消息区域 ===
        send_group = QGroupBox("发送消息")
        send_layout = QVBoxLayout(send_group)
        
        # 发送设置
        settings_layout = QHBoxLayout()
        
        # 间隔设置
        interval_label = QLabel("发送间隔(毫秒):")
        self.interval_input = QLineEdit("1000")
        self.interval_input.setMaximumWidth(80)
        settings_layout.addWidget(interval_label)
        settings_layout.addWidget(self.interval_input)
        
        settings_layout.addStretch()
        
        # JSON格式化选项
        self.format_json_check = QCheckBox("JSON格式")
        self.format_json_check.setChecked(True)
        settings_layout.addWidget(self.format_json_check)
        
        send_layout.addLayout(settings_layout)
        
        # 消息编辑器
        self.message_edit = QTextEdit()
        self.message_edit.setPlaceholderText("在此输入要发送的消息...\n可以是单个JSON对象或JSON数组(批量发送)")
        # 默认发送数据示例
        default_data = [
            {
              "type": "comment",
              "nickname": "用户1",
              "content": "你好啊"
            }
        ]
        self.message_edit.setText(json.dumps(default_data, indent=2, ensure_ascii=False))
        send_layout.addWidget(self.message_edit)
        
        # 发送按钮区域
        send_btn_layout = QHBoxLayout()
        
        # 状态显示
        self.send_status = QLabel("")
        send_btn_layout.addWidget(self.send_status)
        send_btn_layout.addStretch()
        
        # 停止按钮
        self.stop_btn = QPushButton("停止发送")
        self.stop_btn.clicked.connect(self.stop_sending)
        self.stop_btn.setEnabled(False)
        send_btn_layout.addWidget(self.stop_btn)
        
        # 发送按钮
        self.send_btn = QPushButton("发送消息")
        self.send_btn.clicked.connect(self.send_message)
        send_btn_layout.addWidget(self.send_btn)
        send_layout.addLayout(send_btn_layout)
        
        # 添加到主布局
        main_layout.addWidget(send_group)
        
    def init_websocket(self):
        # 创建并启动WebSocket客户端线程
        self.ws_thread = WebSocketThread()
        self.ws_thread.connectionChanged.connect(self.update_connection_status)
        self.ws_thread.start()
    
    def toggle_connection(self):
        if not self.ws_thread:
            return
            
        if self.ws_thread.ws and self.ws_thread.ws.sock and self.ws_thread.ws.sock.connected:
            # 断开连接
            self.ws_thread.stop()
            self.update_connection_status(False, "已断开连接")
            self.connect_btn.setText("连接")
        else:
            # 连接到服务器
            host = self.host_input.text().strip()
            port = self.port_input.text().strip()
            
            try:
                port = int(port)
                if port <= 0 or port > 65535:
                    raise ValueError("端口范围应为1-65535")
            except ValueError as e:
                QMessageBox.warning(self, "输入错误", f"无效的端口: {str(e)}")
                return
                
            self.ws_thread.set_server(host, port)
            self.ws_thread.connect_websocket()
            self.connect_btn.setText("断开")
    
    def update_connection_status(self, connected, message):
        if connected:
            self.status_text.setText("已连接")
            self.status_text.setStyleSheet("color: green; font-weight: bold;")
            self.connect_btn.setText("断开")
        else:
            self.status_text.setText(message)
            self.status_text.setStyleSheet("color: red;")
            self.connect_btn.setText("连接")
    
    def send_message(self):
        if not self.ws_thread:
            QMessageBox.warning(self, "发送失败", "WebSocket客户端未初始化")
            return
            
        # 获取消息内容
        message = self.message_edit.toPlainText().strip()
        if not message:
            QMessageBox.warning(self, "发送失败", "消息内容不能为空")
            return
        
        # 检查是否为JSON格式
        if self.format_json_check.isChecked():
            try:
                # 验证JSON格式
                data = json.loads(message)
                
                # 根据数据类型进行处理
                if isinstance(data, list):
                    if not data:
                        QMessageBox.warning(self, "发送失败", "数组不能为空")
                        return
                        
                    # 获取并验证间隔时间
                    try:
                        interval = int(self.interval_input.text())
                        if interval < 100:  # 最小间隔100毫秒
                            interval = 100
                            self.interval_input.setText("100")
                    except ValueError:
                        QMessageBox.warning(self, "参数错误", "请输入有效的间隔时间(毫秒)")
                        return
                    
                    # 准备批量发送
                    self.data_to_send = data
                    self.current_data_index = 0
                    self.is_sending = True
                    
                    # 更新UI状态
                    self.send_btn.setEnabled(False)
                    self.stop_btn.setEnabled(True)
                    
                    # 开始定时发送
                    self.send_status.setText(f"开始批量发送 (共{len(self.data_to_send)}条, 间隔{interval}毫秒)")
                    
                    # 立即发送第一条
                    self.send_next_data()
                    # 启动定时器发送剩余数据
                    self.send_timer.start(interval)
                    return
                
                # 单个对象直接发送
                message = json.dumps(data)
            except json.JSONDecodeError as e:
                QMessageBox.warning(self, "JSON格式错误", f"消息不是有效的JSON格式: {str(e)}")
                return
        
        # 发送单条消息
        if self.ws_thread.send_message(message):
            self.send_status.setText("消息已发送")
            QTimer.singleShot(3000, lambda: self.send_status.setText(""))
        else:
            QMessageBox.warning(self, "发送失败", "消息发送失败，WebSocket可能未连接")
    
    def stop_sending(self):
        self.is_sending = False
        self.send_timer.stop()
        self.send_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        self.send_status.setText("已停止发送")
        QTimer.singleShot(3000, lambda: self.send_status.setText(""))
    
    def send_next_data(self):
        if not self.is_sending or self.current_data_index >= len(self.data_to_send):
            if self.is_sending:
                self.stop_sending()
                self.send_status.setText("所有数据发送完成")
                QTimer.singleShot(3000, lambda: self.send_status.setText(""))
            return
        
        try:
            # 获取当前数据
            data = self.data_to_send[self.current_data_index]
            self.current_data_index += 1
            
            # 更新状态
            self.send_status.setText(f"正在发送 ({self.current_data_index}/{len(self.data_to_send)})")
            
            # 发送数据
            success = self.ws_thread.send_message(data)
            
            # 检查发送结果
            if not success:
                self.send_status.setText(f"发送失败 ({self.current_data_index}/{len(self.data_to_send)})")
                # 如果发送失败，停止发送
                self.stop_sending()
        except Exception as e:
            self.send_status.setText(f"发送出错: {str(e)}")
            self.stop_sending()
    
    def closeEvent(self, event):
        # 停止WebSocket线程
        if self.ws_thread:
            self.ws_thread.stop()
            self.ws_thread.wait()  # 等待线程结束
        
        event.accept()


if __name__ == "__main__":
    # 设置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_()) 